####################################
#
#
# EXECUTABLE PATH:
# * configure these paths:
#
# these paths are for Phylo...

CLUSTAL = "/Users/victor/Applications/clustalw2"
PRANK = "/Users/victor/Applications/prank.osx_10411_ppc.081202/prank"
AMAP = "/Users/victor/Applications/amap-align/align/amap"
MUSCLE = "/Users/victor/Applications/muscle3.7/muscle"

PROTTEST = "/Users/victor/Applications/ProtTest1.4.mac/scripts/runProtTest"

PHYML = "/common/bin/phyml"
RAXML = "/Users/victor/Applications/RAxML-7.0.4/raxmlHPC"

LAZARUS = "python /usr/bin/lazarus_batch"
SEQCONVERTER = "perl /Users/victor/Applications/seqConverter.pl"
TRANSLATETREES = "python /Users/victor/Applications/PhyloWidgets/translate_tree.py"

MPIDISPATCH = "/Users/victor/Applications/mpi_dispatch/mpi_dispatch"
USE_MPI = False

######################################################
#
# imports and globals. . .
#

import os
import re
import sys
from argParser import ArgParser

argParser = ArgParser(sys.argv)

ERG_SEQS = argParser.getArg("--in") # a filepath to a FASTA file with unaligned sequences
OUT_DIR = argParser.getArg("--out") # a directory where results will be dumped
MODEL = argParser.getOptionalArg("--model")
USE_PINVAR = argParser.getOptionalArg("--usepinvar")
if USE_PINVAR == "True":
    USE_PINVAR = True
USE_GAMMA = argParser.getOptionalArg("--usegamma")
if USE_GAMMA == "True":
    USE_GAMMA = True

alignment_algorithms = ["CLUSTAL", "PRANK", "AMAP", "MUSCLE"]
aligned_fasta = []  # an array of FASTA alignment filepaths
aligned_phy = []    # an array of PHYLIP alignment filepaths
treefiles = []      # an array of treefiles

#####################################################
def Setup():
    if os.path.exists( ERG_SEQS ) == False:
        print "I can't find your sequences at ", ERG_SEQS
        exit(1)
    if os.path.exists( OUT_DIR ) == False:
        print "Your output directory doesn't exist, I'll try to create it."
        if os.path.exists( OUT_DIR ) == False:
            print "Yikes, I can't create your output directory at: ", OUT_DIR
            exit(1)

def Run_mpidispatch(job_array):
    fout = open(OUT_DIR + "/.mpijobs.sh", "w")
    for j in job_array:
        fout.write(j + "\n")
    fout.close()
    if USE_MPI:
        command = "mpirun -n 6 --hostfile ~/.bhosts " + MPIDISPATCH + " " + OUT_DIR + "/.mpijobs.sh"
        print command
    else:
        command = "source " + OUT_DIR + "/.mpijobs.sh"
    os.system(command)

def Get_clustal_string(f):
    return [Get_clustal_string_default(f, OUT_DIR + "/clustal.fasta", 20.0, 2.0), OUT_DIR + "/clustal.fasta"]

def Get_clustal_string_default(input_sequence_file, output_alignment_filename, gap_opening, gap_extension):
    command = CLUSTAL
    command = command + " -INFILE=" + input_sequence_file
    command = command + " -OUTFILE=" + output_alignment_filename
    command = command + " -PWMATRIX=GONNET"
    command = command + " -GAPOPEN=" + gap_opening.__str__()
    command = command + " -TYPE=PROTEIN"
    command = command + " -OUTPUT=FASTA" 
    return command

def Get_prank_string(f):
    command = PRANK
    command += " -d=" + f
    return [command, OUT_DIR + "/output.2.fas"]

def Get_amap_string(f):
    command = AMAP
    command += " " + f
    command += " > amap.fasta"
    return [command, OUT_DIR + "/amap.fasta"]

def Get_muscle_string(f):
    command = MUSCLE
    command += " -in " + f
    command += " -out muscle.fasta"
    return [command, OUT_DIR + "/muscle.fasta"]

def Align():
    job_array = []
    for a in alignment_algorithms:
        if a == "CLUSTAL":
            x = Get_clustal_string(ERG_SEQS)
            job_array.append(x[0])
            aligned_fasta.append(x[1])
        if a == "PRANK":
            x = Get_prank_string(ERG_SEQS)
            job_array.append(x[0])
            aligned_fasta.append(x[1])
        if a == "AMAP":
            x = Get_amap_string(ERG_SEQS)
            job_array.append(x[0])
            aligned_fasta.append(x[1])
        if a == "MUSCLE":
            x = Get_muscle_string(ERG_SEQS)
            job_array.append(x[0])
            aligned_fasta.append(x[1])
    Run_mpidispatch( job_array )

def Create_phy_versions():
    for i in aligned_fasta:
        command = SEQCONVERTER
        command += " -d" + i
        command += " -if"
        command += " -ope"
        command += " -ra"
        os.system(command)
        n = re.sub(".fasta", ".phylip", i)
        aligned_phy.append(n)

def writePhymlScript(filename, model, useinvar, usegamma):
    ##################################################
    # Copied From the PHYML command-line manual:
    # Type 12 parameters for amino-acids sequences :
    # 1. sequences file 
    # 2. data type 
    # 3. sequence format 
    # 4. nb data sets 
    # 5. nb bootstrapped data sets 
    # 6. substitution model 
    # 7. prop. invariable sites 
    # 8. nb categories 
    # 9. gamma parameter 
    # 10. starting tree 
    # 11. optimise topology 
    # 12. optimise branch lengths and rate parameters
    # Example :
    #./phyml seqs2 1 i 1 0 JTT 0.0 4 1.0 BIONJ n n
    f = ""
    f += PHYML
    f += " " + filename #1.
    f += " 1" #2: amino acids
    f += " i" #3: interleaved data format
    f += " 1" #4. 
    f += " 0" #5.
    f += " " + model #6.
    if useinvar:
        f += " e" #7. use the ML proportion of invariant sites
    else:
        f += " 0.0" #7.
    if usegamma:
        f += " 4" #8. use 4 gamma categories
    else:
        f += " 1" #8.
    f += " e" #9. use the ML alpha for gamma
    f += " BIONJ" #10.
    f += " y" #11.
    f += " n" #12.
    return f  

def Find_ml_tree():
    job_array = []
    for i in aligned_phy:
        f = writePhymlScript(i, MODEL, USE_PINVAR, USE_GAMMA)
        job_array.append(f)
        treefiles.append( i + "_phyml_tree.txt" )
    Run_mpidispatch( job_array )

#def Run_lazarus():
#    job_array = []
#    for i in aligned_phy
#    #python /common/share/lazarus/lazarus.py --codeml --outputdir /Users/victor/Documents/fungus_project/TREATMENT_ALPHA/LAZARUS --alignment ./sequences.clustal.fasta --tree ./elided.alignments.phylip_phyml_tree.txt --model /common/share/paml4/dat/wag.dat --asrv 4 --gapcorrect --outgroup [T_annulata_16,M_brevicollis_16,C_neoformans_16,S_purpuratus_16,C_parvum_16]
    
# MAIN
Setup()
Align()
Create_phy_versions()
Find_ml_tree()